import simplejson as json
import os
import sys
from xml.etree.ElementTree import ElementTree
from sets import Set
import httplib, urllib

def main():
    extractSponsorTrainingData('./canada_industry', 100, './sponsor_training');
    #extractSponsors('./canada_industry', 10, './no_stopwords_removed');
    
def test(q):
    sponsor = q;
    q = getReconTrialSponsorQuery(sponsor);
    print q;
    data = postToRecon(q, 0, 10);
    print data;
    if(len(data)>0):
        disam_obj = decide(data, sponsor);
        print disam_obj;
        pretty = json.dumps(disam_obj, separators=(',',':'));
    else :
        pretty = sponsor+":none";
    print pretty;

def checkExcluded(d):
    ex = Set(['/biology/organism_classification',
              '/biology/genome',
              '/biology/gene',
              '/biology/gene_group',
              '/biology/gene_ontology_group',
              '/biology/protein',
              '/book/written_work',
              '/book/book',
              '/broadcast/producer',
              '/business/job_title',
              '/business/industry',
              '/computer/programming_language_designer',
              '/cvg/computer_videogame',
              '/dining/cuisine',
              '/education/field_of_study',
              '/fictional_universe/fictional_character',
              '/film/film',
              '/location/citytown',
              '/location/statistical_region',
              '/location/dated_location',
              '/location/country',
              '/medicine/disease',
              '/medicine/medical_treatment',
              '/music/artist',
              '/music/album',
              '/music/release',
              '/music/song',
              '/people/cause_of_death',
              '/people/person',
              '/people/profession',
              '/time/time_zone',
              '/tv/tv_program'
              ]);
    for t in  d['type']:
        if(t in ex):
            return True;
    return False;

def getAllTypes(reconresponse):
    types = Set([]);
    for d in reconresponse:
        for t in  d['type']:
            types.add(t);
    return types;

def decide(data, qstring):
    typefilter = False;
    max_score = 0;
    high_score_data = data[0];
    for d in data:
    #find the highest scoring d;
        if(d["score"]>max_score):
            max_score = d["score"];
            high_score_data = d;
        #if the recon service says its good, then its good
        if(d['match']=='true'):
            d["choice"] = "recon_match";
            return d;
        #if it contains any of these types, then keep it
        if(typefilter):
            for t in  d['type']:
                if(t=="/medicine/medical_trial_sponsor"):
                    d["choice"] = "trial_sponsor";
                    return d;
                if(t=="/base/bioventurist/science_or_technology_company"):
                    d["choice"] = "bvr_sci_tech";
                    return d;
                if(t=="/medicine/hospital"):
                    d["choice"] = "hospital";
                    return d;
                if(t=="/medicine/cancer_center_constituent"):
                    d["choice"] = "cancer_center";
                    return d;
                if(t=="/base/bioventurist/organization"):
                    d["choice"] = "bvr_org";
                    return d;
                if(t=="/venture_capital/venture_funded_company"):
                    d["choice"] = "venture_funded";
                    return d;
        #see if this data contains an excluded type
  #      excluded = checkExcluded(d);
        #if not excluded and the string matches exactly, keep it
  #      if(excluded == False):
  #          for name in d['name']:
  #              if(name==qstring):
  #                  d["choice"] = "string_match";
  #                  return d;
    #if nothing else found, return one wiht the highest score
    high_score_data["choice"] = "highscore";
    return high_score_data;

def getReconTrialSponsorQuery(name):
    q = "{ \"/type/object/name\":\""+name+"\" , ";
    q+= "\"/type/object/type\":\"/medicine/medical_trial_sponsor\",";
    q+= "\"/type/object/type\":\"/bioventurist/science_or_technology_company\",";
    q+= "\"/type/object/type\":\"/bioventurist/organization\",";
    q+= "\"/type/object/type\":\"/business/employer\","; 
    q+= "\"/type/object/type\":\"/business/company\",";
    q+= "\"/type/object/type\":\"/education/university\",";
    q+= "\"/type/object/type\":\"/organization/organization\", ";
    q+= "\"/type/object/type\":\"/common/topic\"";
    q+= "}";
    return q;

#see http://data.labs.freebase.com/recon/
def postToRecon(q, start, limit):
    url = "data.labs.freebase.com";
    headers = {"Content-type": "application/x-www-form-urlencoded","Accept": "text/json"};
    params = urllib.urlencode({'q':q,'start':start,'limit':limit}); 
    conn = httplib.HTTPConnection(url);
    conn.request("POST", "/recon/query", params, headers); 
    r1 = conn.getresponse();
    #print r1.status, r1.reason;
    data1 = r1.read();
    obj = json.loads(data1);
    conn.close();
    return obj;

def removeSponsorStopwords(text):
    check = text.split();
    ex = Set(['Pharmaceuticals',
              'Incorporated', 
              'Corporation',
              'Inc.',
              'Inc',
              'Therapeutics'
              ]);
    out = "";
    for t in  check:
        if(t in ex):
            print "removed "+t;
        else:
            out = out+t+" ";
    return out;

def extractSponsorTrainingData(dir, filelimit, outroot):
    start = 0;
    limit = 10; 
    c = 0;
    sponsors = Set();
    done = False;
    for root, dirs, files in os.walk(dir):
        if done:
            break;
        for name in files:
            if done:
                break;
            c = c+1;
            if c%100==0:
                print c;
            filename = os.path.join(root, name);
            if(filename.endswith(".xml")):
                tree = parseCT(filename);
                for a in getAgencies(tree):
                    qstring = removeSponsorStopwords(a.text);
                    sponsors.add(qstring);
            if(c>=filelimit):
                done = True;
                break;

    disfile = open(outroot+'.csv', 'w');
    reconned = {};
    #get all the data from recon
    alltypes = Set([]);
    for sponsor in sponsors:
        q = getReconTrialSponsorQuery(sponsor);
        data = postToRecon(q, start, limit);
        alltypes.update(getAllTypes(data));
        if(len(data)>0):
            reconned[sponsor] = data
    
    for sponsor in reconned.iterkeys():  
        for disam_obj in reconned[sponsor]: 
            tabbed = sponsor+"\t"+disam_obj["name"].__str__()+"\t"+disam_obj["id"]+"\t"+disam_obj["score"].__str__()+"\t";
            for t in alltypes:
                if(t in disam_obj["type"]):
                    tabbed = tabbed + t + "\t";
                else:
                     tabbed = tabbed + "0\t";
            print tabbed;
            disfile.write(tabbed);
            disfile.write("\n");
    disfile.close();


def extractSponsors(dir, filelimit, outroot):
    start = 0;
    limit = 10; 
    c = 0;
    sponsors = Set();
    done = False;
    for root, dirs, files in os.walk(dir):
        if done:
            break;
        for name in files:
            if done:
                break;
            c = c+1;
            if c%100==0:
                print c;
            filename = os.path.join(root, name);
            if(filename.endswith(".xml")):
                tree = parseCT(filename);
                for a in getAgencies(tree):
                    qstring = removeSponsorStopwords(a.text);
                    sponsors.add(qstring);
            if(c>=filelimit):
                done = True;
                break;
    jsonfile = open(outroot+'.json', 'w');
    disfile = open(outroot+'.csv', 'w');
    file = open(outroot+'sponsors.txt', 'w');
    file.write("Companies That Sponsored a Clinical Trial in Canada\n");
    for sponsor in sponsors:
        print sponsor;
        q = getReconTrialSponsorQuery(sponsor);
  #      print q;
        data = postToRecon(q, start, limit);
  #      print data;
        if(len(data)>0):
            disam_obj = decide(data, sponsor);
            disam_obj["querystring"] = sponsor;
  #      print disam_obj;
            pretty = json.dumps(disam_obj, separators=(',',':'));
            tabbed = disam_obj["querystring"]+"\t"+disam_obj["name"].__str__()+"\t"+disam_obj["choice"]+"\t"+disam_obj["id"]+"\t"+disam_obj["score"].__str__()+"\t"+disam_obj["type"].__str__()+"\t"+disam_obj["match"].__str__();
        else :
            pretty = sponsor+":none";
            tabbed = sponsor+"\tnone\tnone\tnone\t0\tnone\tFALSE\t";
  #      print disam_obj;
        print tabbed;
        disfile.write(tabbed);
        disfile.write("\n");
        jsonfile.write(pretty);
        jsonfile.write('\n');
        file.write(sponsor);
        file.write('\n');
    file.close();
    disfile.close();
    
def parseCT(infile):
    tree = ElementTree();
    tree.parse(infile);
    return tree; 

def getConditions(tree):
    return tree.findall('.//condition');

def getAgencies(tree):
    return tree.findall('.//agency');

if __name__ == "__main__":
  main();